#!/usr/bin/env python3
# coding: utf-8
from pprint import pprint
from logs.logger import logger
from mod.webapi import WebApi, APKEncryptWebApi, H5EncryptWebApi
import time, re


class ApkEncErrorField():
    APKENC_ERROR_TOOL = 1         # 获取工具失败
    APKENC_ERROR_UPLOAD = 2       # 上传文件失败
    APKENC_ERROR_PROJECT = 3      # 获取项目失败
    APKENC_ERROR_START = 4        # 开始加固失败
    APKENC_ERROR_LOOP = 5         # 轮询状态失败
    APKENC_ERROR_ENCRYPT = 6      # 加固失败
    APKENC_ERROR_DOWNLOAD = 7     # 下载文件失败
    APKENC_ERROR_CANCEL = 8       # 加固任务被取消
    APKENC_ERROR_TOOL_USE = 9     # 加固版本匹配错误
    APKENC_ERROR_POLICY_USE = 9   # 加固策略匹配错误


class ApkEncryptor(ApkEncErrorField):
    def __init__(self, apkin='', apkout='', pid='', toolver=''):
        self.apkin = apkin
        self.apkout = apkout
        self.pid = pid
        self.toolver = toolver
        self.webapi = APKEncryptWebApi()
        return

    def throw_err(self, err, task_id=None):
        name = ''
        for k,v in vars(ApkEncErrorField).items():
            if v == err:
                name = k
                break
        logger.error(f'[-] Raise error {name} code:{err}')

        if task_id:
            msg = self.webapi.delete_task(task_id)
            logger.info(f'[-] delete task_id {task_id} message: {msg}')
        raise err

    def loop(self, task_id, threshold):
        retries = threshold
        while True:
            time.sleep(5)
            status = self.webapi.get_encrypt_status(task_id)
            if not status:
                if retries <= 0:
                    logger.error(f'[-] 查询加固任务状态失败, 加固程序退出！')
                    return False
                retries -= 1
                continue
            # 0 应用处于初始状态，未进入加固
            # 1 加固中
            # 2 加固完成
            # 3 队列中
            # 4 应用已经加固过，无法重复加固
            # -1 加固失败
            # retries恢复，若连续失败次数超过threshold就失败
            retries = threshold
            logger.info(f'[+] 当前加固状态: {status["protectTaskStr"]}')
            if status['protectTaskStr'] == '失败':
                error_type_map = {
                    "33": "请检查文件完整性校验列表文件路径是否合法",
                    "105": "请检查AndroidManifest是否已加密",
                    "108": "请检查apk是否签名",
                }
                err_msg = '未知加固异常'
                if status["errMsg"] and re.match("Traceback", status["errMsg"]) and re.match("self._download", status["errMsg"]):
                    err_msg = '网络波动导致加固服务器无法下载待加固文件'
                elif status["errMsg"]:
                    err_msg = error_type_map.get(status["errMsg"], '未知加固异常')
                logger.error(f'[-] {err_msg}, 加固程序退出！')
                return False
            elif status['protectTaskStr'] in ['配置中', '已取消']:
                logger.error(f'[-] 加固任务被取消, 加固程序退出！')
                return False
            elif status['protectTaskStr'] == '已完成':
                return True

    def start(self, task_id, pid, toolver):
        # 每次使用cli都是上传一个加固任务，然后开始这个任务
        result = self.webapi.start_encrypt(task_id, pid, toolver)
        if not result:
            return False

        logger.info(f'[+] 创建加固任务成功')
        return True

    def run(self):
        # 获取工具列表
        tool = self.webapi.get_toolVersion()
        if not tool:
            logger.error(f'[-] 获取加固版本list数据失败, 加固程序退出！')
            return
        show_tool = [item['name'] for item in tool]
        logger.info(f'[+] 可使用的加固版本list: {show_tool}')

        # 获取策略列表
        policy_infos = self.webapi.get_encrypt_policys()
        if not policy_infos:
            logger.error(f'[-] 获取加固策略list数据失败, 加固程序退出！')
            return
        show_policy_list = [item['policy_id'] for item in policy_infos]
        logger.info(f'[+] 可选择的加固策略list: {show_policy_list}')

        # 匹配加固版本，加固策略
        if not self.pid:
            self.pid = policy_infos[0]['policy_id']
            logger.info(f'[+] 命令行缺少策略id, 默认使用全局加固策略id: {self.pid}进行加固。')
        if not self.toolver:
            self.toolver = show_tool[0]
            logger.info(f'[+] 命令行缺少指定加固版本, 默认使用最新加固版本: {self.toolver}进行加固。')

        if int(self.pid) not in [p['policy_id'] for p in policy_infos]:
            logger.error(f'[-] 错误的加固策略id: {self.pid}, 加固程序退出！')
            return
        if self.toolver not in show_tool:
            logger.error(f'[-] 错误的加固版本: {self.toolver}, 加固程序退出！')
            return

        # 上传文件
        task_id = self.webapi.upload_file(self.apkin)
        if not task_id:
            return
        logger.info(f'[+] 上传文件成功, task_id:{task_id}')

        # 开始加固
        ret = self.start(task_id, self.pid, {item['name']: item['value'] for item in tool}[self.toolver])
        if not ret:
            logger.error(f'[-] 创建加固任务失败, 加固程序退出！')
            return

        # 轮询状态
        task_ret = self.loop(task_id, threshold=6)
        if not task_ret:
            return
        logger.info(f'[+] 文件加固成功！开始下载加固文件')

        # 下载文件
        r = self.webapi.down_file(task_id, self.apkout)
        if not r:
            logger.error(f'[-] 加固包下载失败，请检查存放文件目录是否存在')
        logger.info(f'[+] 加固程序执行成功, 加固包路径:{self.apkout}')

        return True


class H5Encryptor(ApkEncErrorField):
    def __init__(self, apkin='', apkout='', pid=''):
        self.apkin = apkin
        self.apkout = apkout
        self.pid = pid
        self.webapi = H5EncryptWebApi()
        return

    def throw_err(self, err, task_id=None):
        name = ''
        for k,v in vars(ApkEncErrorField).items():
            if v == err:
                name = k
                break
        logger.error(f'[-] Raise error {name} code:{err}')

        if task_id:
            msg = self.webapi.delete_task(task_id)
            logger.info(f'[-] delete task_id {task_id} message: {msg}')
        raise err

    def loop(self, task_id, threshold):
        retries = threshold
        while True:
            time.sleep(5)
            status = self.webapi.get_encrypt_status(task_id)
            if not status:
                if retries <= 0:
                    logger.error(f'[-] 查询加固任务状态失败, 加固程序退出！')
                    return False
                retries -= 1
                continue
            retries = threshold
            logger.info(f'[+] 当前加固状态: {status["protectTaskStr"]}')
            if status['protectTaskStr'] == '失败':
                err_msg = '未知加固异常'
                if status.get("errMsg"):
                    err_msg = status["errMsg"]
                logger.error(f'[-] {err_msg}, 加固程序退出！')
                return False
            elif status['protectTaskStr'] in ['配置中', '已取消']:
                logger.error(f'[-] 加固任务被取消, 加固程序退出！')
                return False
            elif status['protectTaskStr'] == '已完成':
                return True

    def start(self, task_id, pid):
        # 每次使用cli都是上传一个加固任务，然后开始这个任务
        result = self.webapi.start_encrypt(task_id, pid)
        if not result:
            return False

        logger.info(f'[+] 创建加固任务成功')
        return True

    def run(self):
        # 获取策略列表
        policy_infos = self.webapi.get_encrypt_policys()
        if not policy_infos:
            logger.error(f'[-] 获取加固策略list数据失败, 加固程序退出！')
            return
        show_policy_list = [item['rule_id'] for item in policy_infos]
        logger.info(f'[+] 可选择的加固策略list: {show_policy_list}')

        # 匹配加固版本，加固策略
        if not self.pid:
            self.pid = policy_infos[-1]['rule_id']
            logger.info(f'[+] 命令行缺少策略id, 默认使用全局加固策略id: {self.pid}进行加固。')
        if self.pid not in [p['rule_id'] for p in policy_infos]:
            logger.error(f'[-] 错误的加固策略id: {self.pid}, 加固程序退出！')
            return

        # 上传文件
        task_id = self.webapi.upload_file(self.apkin)
        if not task_id:
            return
        logger.info(f'[+] 上传文件成功, task_id:{task_id}')

        # 开始加固
        ret = self.start(task_id, self.pid)
        if not ret:
            logger.error(f'[-] 创建加固任务失败, 加固程序退出！')
            return

        # 轮询状态
        task_ret = self.loop(task_id, threshold=6)
        if not task_ret:
            return
        logger.info(f'[+] 文件加固成功！')

        # 下载文件
        logger.info(f'[+] 开始下载加固包')
        r = self.webapi.down_file(task_id, self.apkout)
        if not r:
            logger.error(f'[-] 加固包下载失败，请检查存放文件目录是否存在')
        logger.info(f'[+] 加固程序执行成功, 加固包路径:{self.apkout}')
        return True